# -*- coding:utf-8 -*-
import collections
from elasticsearch import Elasticsearch
import time

STCPostCmntPair = collections.namedtuple("STCPostCmntPair",
                                         ("id",
                                          "post",
                                          "cmnt",
                                          "score"))


def parse_query_result(hit):
    post = hit["_source"]["post"]
    cmnt = hit['_source']["cmnt"]
    score = hit['_score']
    doc_id = hit['_id']
    return STCPostCmntPair(id=doc_id,
                           post=post,
                           cmnt=cmnt,
                           score=score)


class ElasticSearchHelper(object):
    def __init__(self):
        self._client = Elasticsearch()
        # self._client = Elasticsearch(hosts="219.224.169.14", port=9200)
        pass

    def query(self, query, size=100):
        res = self._client.search(index="stc-index",
                                  doc_type='atec',
                                  size=size,
                                  body={"query": {"match": {'cmnt': query}}})
        print('-' * 100)
        ret_items = []
        for hit in res['hits']['hits']:
            ret_items += [parse_query_result(hit)]
        return ret_items


def test():
    es_helper = ElasticSearchHelper()
    start = time.time()
    for doc in es_helper.query("今天天气真好"):
        print(doc)
    print("time cost: {}".format(time.time() - start))


def generate_stc_training_data(stc_data_path, saved_path):
    queries = set([])
    with open(stc_data_path, 'r', encoding="utf-8") as fr:
        for line in fr:
            ss = line.rstrip().split("\t")
            queries.add(ss[0])

    print("query size={}".format(len(queries)))
    es_helper = ElasticSearchHelper()
    with open(saved_path, 'w', encoding="utf-8") as fw:
        for idx, query in enumerate(queries):
            items = es_helper.query(query, 600)
            for item in items[100:]:
                fw.write("{}\t{}\t{}\t{}\n".format(query, item.score, item.post, item.cmnt))
            if (idx + 1) % 100 == 0:
                print("index {} ...".format(idx + 1))

    pass


if __name__ == '__main__':
    test()
    # generate_stc_training_data(sys.argv[1], sys.argv[2])
    pass
